#include "bitmap.h"

#define bm_index(bm, pos) ( (pos) >> (bm)->shift )
#define bm_submask(bm) ( ( 1 << (bm)->shift ) - 1 )
#define bm_subsize(bm) ( 1 << (bm)->shift )
#define bm_subindex(bm, pos) ( (pos) & bm_submask((bm)) )
#define bm_has_child(bm) ( (bm)->shift != 5 )
#define bm_pos(bm, i, j) ( ((i) << (bm)->shift) | ((j) & bm_submask((bm))) )

static inline u8 bm_size_to_shift(u32 size)
{
	u8 shift = 0;

	if (size > 0)
  		--size;
	do {
		size >>= 5;
		shift += 5;
	} while (size >= 32);

	return shift;
};

#ifdef __KERNEL__

#include <linux/slab.h>
#define mem_alloc(size) kmalloc((size), GFP_KERNEL)
#define mem_free(p) kfree((p))

#define BM_ERROR(fmt, args...) printk(KERN_ERR "bitmap:%s:%d " fmt "\n", \
    __FUNCTION__, __LINE__, ##args)
#define TRACE(fmt, args...) printk(KERN_DEBUG "bitmap:%s:%d " fmt "\n", \
    __FUNCTION__, __LINE__, ##args)

#else

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#define mem_alloc malloc
#define mem_free free

#define BM_ERROR(fmt, args...) fprintf(stderr, "bitmap:%s:%d " fmt "\n", \
    __FUNCTION__, __LINE__, ##args)
#define TRACE(fmt, args...) printf("bitmap:%s:%d " fmt "\n", \
    __FUNCTION__, __LINE__, ##args)
#include <asm-generic/errno-base.h>

#endif

struct bitmap* bm_alloc(u32 size)
{
	struct bitmap* p = mem_alloc(sizeof(struct bitmap));
	if (!p) {
		BM_ERROR("error, no memory");
		return NULL;
	}
	memset(p, 0, sizeof(*p));
	p->size = size;
	p->shift = bm_size_to_shift(size);
	//TRACE("bm_alloc %p, size %u", p, p->size);
	return p;
}

void bm_free(struct bitmap* bm)
{
	if (bm_has_child(bm)) {
		int i;
		for (i = 0; i < 32; ++i) {
			if (bm->pbm[i] != 0)
				bm_free(bm->pbm[i]);
		}
	}
	mem_free(bm);
	//TRACE("bm_free %p", bm);
}

int bm_getb(struct bitmap* bm, u32 pos)
{
	u32 i, j;

	if (pos >= bm->size) {
		BM_ERROR("get pos (%u) out range of size (%u)", pos, bm->size);
		return 0;
	}

	i = bm_index(bm, pos);
	j = bm_subindex(bm, pos);
	//TRACE("i = %u, j = %u, pos = %u", i, j, pos);

	if (bm_has_child(bm)) {
		return (bm->pbm[i]) ? bm_getb(bm->pbm[i], j) : 0;
	}
	else {
		//TRACE("get bits = %u", bm->bits[i]);
		return u32_isset(bm->bits[i], j);
	}
}

int bm_setb(struct bitmap* bm, u32 pos, int val)
{
	u32 i, j;

	if (pos >= bm->size) {
		BM_ERROR("set pos (%u) out range of size (%u)", pos, bm->size);
		return -EINVAL;
	}
	
	i = bm_index(bm, pos);
	j = bm_subindex(bm, pos);
	//TRACE("i = %u, j = %u, pos = %u", i, j, pos);

	if (bm_has_child(bm)) {
		if (bm->pbm[i] == NULL && val != 0) {
			bm->pbm[i] = bm_alloc(bm_subsize(bm));
			if (!bm->pbm[i])
				return -ENOMEM;
		}
		if (bm->pbm[i])
			return bm_setb(bm->pbm[i], j, val);
	}
	else {
		//TRACE("before bits = %u", bm->bits[i]);
		bm->bits[i] = u32_setb(bm->bits[i], j, val);
		//TRACE("after bits = %u", bm->bits[i]);
	}
	return 0;
}

int bm_nextb(struct bitmap* bm, u32* pos)
{
	u32 i, j, maxi;
	if (*pos >= bm->size)
		return 0;

	i = bm_index(bm, *pos);
	j = bm_subindex(bm, *pos);
	maxi = bm_index(bm, bm->size - 1) + 1;

	if (bm_has_child(bm)) {
		for(;i < maxi; ++i, j = 0) {
			if (bm->pbm[i] == NULL)
				continue;
			if (bm_nextb(bm->pbm[i], &j)) {
				*pos = bm_pos(bm, i, j);
				//TRACE("pos = %u, i = %u, j = %u", *pos, i, j);
				return 1;
			}
		}
		return 0;
	}
	else {
		for(;i < maxi; ++i, j = 0) {
			if (bm->bits[i] == 0)
				continue;
			for(; j < 32; ++j) {
				if (u32_isset(bm->bits[i], j)) {
					*pos = bm_pos(bm, i, j);
					//TRACE("pos = %u, i = %u, j = %u", *pos, i, j);
					return 1;
				}
			}
		}
		return 0;
	}

}

size_t bm_enc_size(struct bitmap* bm)
{
	u32 meta_bm = 0;
	int i;
	int size = 0;
	for (i = 0; i < 32; ++i) {
		if (bm_has_child(bm) && bm->pbm[i] != NULL) {
			int s = bm_enc_size(bm->pbm[i]);
			if (s == 0)
				continue;
			meta_bm = u32_setb(meta_bm, i, 1);
			size += s;
		}
		else if (!bm_has_child(bm) && bm->bits[i] != 0) {
			meta_bm = u32_setb(meta_bm, i, 1);
			++size;
		}
	}

	if (meta_bm == 0)
		return 0;
	else
		return size + 1;
}

size_t bm_enc_buf(struct bitmap* bm, u32* buf, size_t size)
{
	size_t off = 1;
	u32 meta = 0;
	int i;

	if (size == 0)
		return 0;
	
	for (i = 0; i < 32; ++i) {
		if (off >= size)
			break;
		if (bm_has_child(bm) && bm->pbm[i] != NULL) {
			int ret = bm_enc_buf(bm->pbm[i], buf + off, size - off);
			if (ret != 0) {
				off += ret;
				meta = u32_setb(meta, i, 1);
			}
		}
		else if (!bm_has_child(bm) && bm->bits[i] != 0) {
			buf[off++] = bm->bits[i];
			meta = u32_setb(meta, i, 1);
		}
	}
	
	if (off == 1)
		return 0;
	buf[0] = meta;
	return off;
}

size_t bm_dec_buf(struct bitmap* bm, const u32* buf, size_t size)
{
	int off;
	u32 meta;
	int i;

	if (size == 0)
		return 0;
	
	meta = buf[0];
	off = 1;
	for (i = 0; i < 32; ++i) {
		if (off >= size)
			break;
		if (u32_isset(meta, i)) {
			if (bm_has_child(bm)) {
				if (bm->pbm[i] == NULL)
					bm->pbm[i] = bm_alloc(bm_subsize(bm));
				off += bm_dec_buf(bm->pbm[i], buf + off, size - off);
			}
			else {
				bm->bits[i] = buf[off++];
			}
		}
	}
	return off;
}

u32 bm_count(struct bitmap* bm)
{
	u32 pos = 0;
	u32 count = 0;
	while (bm_nextb(bm, &pos)) {
		++count;
		++pos;
	}
	return count;
}
